{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bb63114",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "原始数据形状: (29067, 168)\n",
      "NaN值数量: 976826\n",
      "预处理后数据形状: (29067, 168)\n",
      "预处理后NaN值数量: 976826\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:276: UserWarning: shlt have very rare categories, it is a good idea to group these, or set the min_data_in_leaf parameter to prevent lightgbm from outputting 0.0 probabilities.\n",
      "  warn(\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n",
      "/media/disk/02drive/13hias/miniconda3/envs/jyx-py39/lib/python3.9/site-packages/miceforest/imputation_kernel.py:867: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`\n",
      "  self.candidate_preds[variable][assign_col_index] = candidate_preds\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import miceforest as mf\n",
    "\n",
    "# 读取原始数据\n",
    "df = pd.read_pickle(r\"../data/data_new.pkl\")\n",
    "print(\"原始数据形状:\", df.shape)\n",
    "\n",
    "# 重置索引，确保索引是连续整数\n",
    "df_reset = df.reset_index(drop=True)\n",
    "\n",
    "print(f\"NaN值数量: {df_reset.isnull().sum().sum()}\")\n",
    "\n",
    "for col in df_reset.select_dtypes(include=[\"object\"]).columns:\n",
    "    df_reset[col] = df_reset[col].astype(\"category\")\n",
    "df_reset = df_reset.replace([np.inf, -np.inf], np.nan)\n",
    "\n",
    "print(f\"预处理后数据形状: {df_reset.shape}\")\n",
    "print(f\"预处理后NaN值数量: {df_reset.isnull().sum().sum()}\")\n",
    "\n",
    "# 使用 miceforest 进行多重插补\n",
    "kernel = mf.ImputationKernel(\n",
    "    df_reset,\n",
    "    num_datasets=5,                 # 生成5个插补数据集\n",
    "    random_state=42\n",
    ")\n",
    "kernel.mice(3, min_sum_hessian_in_leaf=5, min_data_in_leaf=10)  # 进行3轮插补，设置最小叶子节点数据量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eea50caf",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'kernel' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m imputed_dfs \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m5\u001b[39m):\n\u001b[0;32m----> 4\u001b[0m     df_imp \u001b[38;5;241m=\u001b[39m \u001b[43mkernel\u001b[49m\u001b[38;5;241m.\u001b[39mcomplete_data(dataset\u001b[38;5;241m=\u001b[39mi)\n\u001b[1;32m      5\u001b[0m     \u001b[38;5;66;03m# 插入 impute_id 和 subject_id\u001b[39;00m\n\u001b[1;32m      6\u001b[0m     df_imp\u001b[38;5;241m.\u001b[39minsert(\u001b[38;5;241m0\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mimpute_id\u001b[39m\u001b[38;5;124m'\u001b[39m, i)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'kernel' is not defined"
     ]
    }
   ],
   "source": [
    "# 保存和拼接所有插补结果\n",
    "imputed_dfs = []\n",
    "for i in range(5):\n",
    "    df_imp = kernel.complete_data(dataset=i)\n",
    "    # 插入 impute_id 和 subject_id\n",
    "    df_imp.insert(0, 'impute_id', i)\n",
    "    df_imp.insert(1, 'subject_id', range(len(df_imp)))\n",
    "    # 保存单个插补数据集\n",
    "    df_imp.to_pickle(f\"./data/imputed/imputed_dataset_{i}.pkl\")\n",
    "    imputed_dfs.append(df_imp)\n",
    "\n",
    "# # 拼接成一个大表\n",
    "# df_imputed = pd.concat(imputed_dfs, axis=0, ignore_index=True)\n",
    "# df_imputed.to_pickle(r\"../data/imputed/imputed_dataset_all.pkl\")\n",
    "\n",
    "# # 打印检查信息\n",
    "# print(\"\\n拼接后的 df_imputed：\")\n",
    "# print(df_imputed.shape)\n",
    "# print(df_imputed.index)\n",
    "# print(df_imputed.columns)\n",
    "# print(df_imputed.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80454dbb",
   "metadata": {},
   "source": [
    "单次插补，使用sklearn的bayesianridge，作为敏感性分析素材"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1df067f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.impute import SimpleImputer, IterativeImputer\n",
    "from sklearn.compose import ColumnTransformer\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.linear_model import BayesianRidge\n",
    "\n",
    "\n",
    "# 读取数据\n",
    "df = pd.read_pickle(r\"../data/data_new.pkl\")\n",
    "\n",
    "# 假设 df 是你的数据集\n",
    "numeric_cols = df.select_dtypes(include=[\"float64\", \"int64\"]).columns\n",
    "categorical_cols = df.select_dtypes(include=[\"category\", \"object\"]).columns\n",
    "\n",
    "# 连续型插补\n",
    "iter_imp = IterativeImputer(estimator=BayesianRidge(), random_state=0)\n",
    "df[numeric_cols] = iter_imp.fit_transform(df[numeric_cols])\n",
    "\n",
    "# 类别型插补\n",
    "cat_imp = SimpleImputer(strategy=\"most_frequent\")\n",
    "df[categorical_cols] = cat_imp.fit_transform(df[categorical_cols])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb94ff0f",
   "metadata": {},
   "source": [
    "分箱"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4d55cc3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    age age_group\n",
      "0  56.0      0-64\n",
      "1  58.0      0-64\n",
      "2  60.0      0-64\n",
      "3  62.0      0-64\n",
      "4  64.0      0-64\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# 读取数据\n",
    "df = pd.read_pickle(r\"../data/imputed/imputed_dataset_all.pkl\")\n",
    "\n",
    "bins = [0, 64, 110]                # 分割点\n",
    "labels = [\"0-64\", \"65-110\"]        # 每个区间的名字\n",
    "df[\"age_group\"] = pd.cut(df[\"age\"], bins=bins, labels=labels, right=True)\n",
    "\n",
    "bins = [-0.1, 6, 8]\n",
    "labels = [\"0-5\",\"6-8\"]\n",
    "df[\"cesd_group\"] = pd.cut(df[\"cesd\"], bins=bins, labels=labels, right=False)\n",
    "\n",
    "print(df[[\"age\", \"age_group\"]].head())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ede4dfb3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "age_group\n",
      "65-110    87285\n",
      "0-64      58050\n",
      "Name: count, dtype: int64\n",
      "cesd_group\n",
      "0-5    139110\n",
      "6-8      4995\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# 统计每一组的数据量\n",
    "print(df[\"age_group\"].value_counts())\n",
    "print(df[\"cesd_group\"].value_counts())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "56720494",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存\n",
    "df.to_pickle(r\"../data/imputed/imputed_dataset_all_binned.pkl\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jyx-py39",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.23"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
